# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import os
import logging
import numpy
import pandas

from absl import app, flags
from concurrent.futures import ProcessPoolExecutor

from coin.base.datetime_util import to_timestamp_int
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog, OrderEvent
from coin.pnl.sim_stat_plot import setup_plt
import coin.proto.coin_order_enums_pb2 as coin_order_pb2
import coin.strategy.mm.tool.archive_base as abase

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


class OrderAcceptLatency(object):
  def __init__(self, trading_date, start_time, end_time, prefix):
    self.nametuple = prefix
    self.trading_date = trading_date
    self.start_time = to_timestamp_int(start_time)
    self.end_time = to_timestamp_int(end_time)
    self._key2subtime = {}
    self._latencies = {}
    self._key2cantime = {}
    self._canlatencies = {}
    self._key2takerfilltime = {}
    self._taker_fill_latencies = {}

  def on_log2(self, timestamp, log):
    if self.start_time <= timestamp <= self.end_time:
      pass
    else:
      return

    pb = OrderGatewayLog()
    try:
      pb.ParseFromString(log)
    except Exception:
      print(log)
      return
    if pb.type == OrderGatewayLog.ORDER_EVENT:
      event = pb.event
      if event.type in [
          OrderEvent.ORDER_SUBMITTED, OrderEvent.ORDER_ACCEPTED, OrderEvent.ORDER_FILLED
      ]:
        key = (event.symbol, event.internal_order_id)
      elif event.type in [OrderEvent.CANCEL_SUBMITTED, OrderEvent.CANCEL_CONFIRMED]:
        key = (event.symbol, event.external_order_id)
      else:
        return
      event_time = int(event.event_time)
      if event.type == OrderEvent.ORDER_SUBMITTED:
        self._key2subtime[key] = event_time
      elif event.type == OrderEvent.ORDER_ACCEPTED:
        if key in self._key2subtime:
          self._latencies[event_time] = event_time - self._key2subtime[key]
      elif event.type == OrderEvent.CANCEL_SUBMITTED:
        if key not in self._key2cantime:
          self._key2cantime[key] = event_time
      elif event.type == OrderEvent.CANCEL_CONFIRMED:
        if key in self._key2cantime:
          self._canlatencies[event_time] = event_time - self._key2cantime[key]
      elif (event.type == OrderEvent.ORDER_FILLED
            and event.fill_type == coin_order_pb2.TAKER_FILL_TYPE):
        if key in self._key2subtime:
          self._taker_fill_latencies[event_time] = event_time - self._key2subtime[key]

  @staticmethod
  def plot_latency(column, latencies, ax1, ax2):
    legends = []
    lats = pandas.Series(latencies) * 1e-6
    ax1.plot(pandas.DatetimeIndex(lats.index), lats, 'r.', markersize=0.5)
    ax1.set_title(column, fontsize=12)
    seri = lats
    for len_unit, col in [(10, 'g'), (100, 'k')]:
      if len(latencies) < len_unit:
        continue

      seri.index = pandas.DatetimeIndex(seri.index.tolist())
      series = seri.sort_index()

      ax1.plot(series.rolling(len_unit).median(), f'{col}.-', linewidth=0.5, markersize=0.5)
      legends.append(f"k={len_unit} rolling avg(ms)")
    ax1.minorticks_on()
    ax1.grid(lw=0.5, which='major')
    ax1.grid(lw=0.2, which='minor')
    ax1.legend([f"{column}(ms)"] + legends, loc="upper right", fontsize=10)

    df1 = pandas.DataFrame(seri.quantile(numpy.linspace(0, 0.5, 11))).reset_index(drop=False)
    df2 = pandas.DataFrame(seri.quantile(numpy.linspace(0.5, 1, 11))).reset_index(drop=False)
    df = pandas.concat([df1, df2], axis=1)
    df.columns = ['qtl', 'lat(ms)', 'qtl', 'lat(ms)']
    df = df.applymap('{:,.2f}'.format)
    tbl = ax2.table(cellText=df.values, loc='center', fontsize=12, colLabels=df.columns)
    tbl.scale(1, 1.4)
    ax2.axis('off')

  def export_log(self):
    setup_plt()
    plt.rcParams["figure.figsize"] = 10, 10
    colnames = ["ordersub_to_ack_lat", "cansub_to_confirm_lat", "sub_to_taker_fill_lat"]
    values = [self._latencies, self._canlatencies, self._taker_fill_latencies]
    plt.suptitle(self.nametuple + self.trading_date.strftime("%Y%m%d"), fontsize=15)

    from matplotlib import gridspec
    gs = gridspec.GridSpec(3, 3)

    for i, (colname, value) in enumerate(zip(colnames, values)):
      ax1 = plt.subplot(gs[i, 0:2])
      ax2 = plt.subplot(gs[i, 2])
      self.plot_latency(colname, value, ax1, ax2)

    os.makedirs(flags.FLAGS.out_dir, exist_ok=True)
    filepath = f'{flags.FLAGS.out_dir}/latency_{self.nametuple}.png'
    print(filepath)
    plt.savefig(filepath)
    plt.close()


def launch(trading_date, baseflags, fsflags, osflags, mea):
  baseflags.market_type, baseflags.exchange = mea.split(":")[:2]
  baseflags.trading_date = trading_date.strftime("%Y%m%d")
  start_time, end_time = abase.get_time_range(baseflags)
  strat = OrderAcceptLatency(trading_date,
                             start_time,
                             end_time,
                             prefix=f"{mea}_{osflags.orderlog_machine}_{osflags.owner}_")
  logging.info("start log analysis")
  print(baseflags, osflags)
  abase.run_from_order_archive_base(baseflags, osflags, on_log=strat.on_log2)
  logging.info("done")
  strat.export_log()
  return strat


def main(_):
  baseflags = abase.get_base_flags()
  fsflags = abase.get_feed_subscription_flags()
  osflags = abase.get_order_subscription_flags()

  trading_dates = abase.get_trading_dates(baseflags.trading_date)
  args = tuple([baseflags, fsflags, osflags, flags.FLAGS.mea])
  strategies = []
  if len(trading_dates) == 1:
    strategies.append(launch(trading_dates[0], *args))
  else:
    futures = []
    with ProcessPoolExecutor(max_workers=16) as executor:
      for trading_date in trading_dates:
        futures.append(executor.submit(launch, trading_date, *args))
      for future in futures:
        strategies.append(future.result())


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  abase.define_order_archive_flags()

  flags.DEFINE_string('mea', None, 'mea name')
  flags.DEFINE_string('out_dir', 'oe_stat', 'out dir')
  flags.DEFINE_bool('detail', False, 'show detail stats')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
